package cn.iocoder.yudao.module.engine.model.context;

import cn.iocoder.yudao.module.engine.enums.AudioSendTaskStateEnum;
import cn.iocoder.yudao.module.engine.enums.EnumActionType;
import cn.iocoder.yudao.module.engine.model.aichat.AIChatFuture;
import cn.iocoder.yudao.module.engine.model.aichat.AIChatReplayMessage;
import cn.iocoder.yudao.module.engine.model.aichat.AIChatReq;
import cn.iocoder.yudao.module.engine.service.core.aichat.splitter.TextSplitter;
import cn.iocoder.yudao.module.resource.api.model.StoryDTO;
import cn.iocoder.yudao.module.resource.api.model.message.AudioDataDTO;
import lombok.AllArgsConstructor;
import lombok.Data;

import java.io.IOException;
import java.time.Instant;
import java.util.*;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.Semaphore;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

@Data
public class AIMessageContext {
    private final AtomicInteger audioReplyCount = new AtomicInteger(0);

    // input
    private final String gameMatchId;
    private final String traceId;
    private final AIChatReq req;
    private final StoryDTO story;

    private final AIChatFuture aiChatFuture;
    private final TextSplitter textSplitter;

    // llm
    private Long requestLLMStartTime = null;
    private Long requestLLMEndTime = null;
    private String fullReplyText = null;
    private final List<String> replyTexts = new ArrayList<>();
    private final Map<String, List<AudioDataDTO>> sentence4AudioData = new HashMap<>();

    // tts
    private final AtomicBoolean textDoneForAudio = new AtomicBoolean(false);
    private final List<AudioTask> audioTaskQueue = Collections.synchronizedList(new ArrayList<>()); // 顺序任务队列
    private final Object audioTaskLock = new Object(); // 用于控制发送的锁

    private String messageId = null;
    private boolean fromLLMCache = false;
    private boolean fromTTSCache = false;

    public Instant getWebsocketStart() {
        return req.getWebsocketStart();
    }

    public int getRequestLLMTimeMs() {
        if (requestLLMStartTime == null || requestLLMEndTime == null) {
            return 0;
        }
        return (int) (requestLLMEndTime - requestLLMStartTime);
    }

    public String getUserText() {
        return aiChatFuture.getUserText();
    }

    public String getChatId() {
        return aiChatFuture.getChatId();
    }

    public int getTtsFuturesSize() {
        return getAiChatFuture().getTtsFuturesSize();
    }

    public Semaphore getSemaphore(EnumActionType actionType) {
        return req.getChatClient().getSemaphore(actionType);
    }

    public void complete(AIMessageContext messageCtx) {
        req.getChatClient().complete(messageCtx);
    }

    public void completeWithError(AIMessageContext messageCtx, Throwable ex) {
        req.getChatClient().completeWithError(messageCtx, ex);
    }

    public void send(AIMessageContext messageCtx, AIChatReplayMessage msg) throws IOException {
        req.getChatClient().send(messageCtx, msg);
    }

    public int remainingSendTaskCount() {
        synchronized (audioTaskQueue) {
            int count = 0;
            for (AudioTask audioTask : audioTaskQueue) {
                if (audioTask.getSendState() != AudioSendTaskStateEnum.DONE) {
                    count++;
                }
            }
            return count;
        }
    }

    public int getTtsNotDoneCount() {
        synchronized (audioTaskQueue) {
            int count = 0;
            for (AudioTask audioTask : audioTaskQueue) {
                if (!audioTask.isTtsDone()) {
                    count++;
                }
            }
            return count;
        }
    }

    public AudioTask getFirstNeedSendTask() {
        synchronized (audioTaskQueue) {
            AudioTask task = null;
            for (AudioTask audioTask : audioTaskQueue) {
                if (audioTask.isTtsDone()) {
                    final AudioSendTaskStateEnum sendState = audioTask.getSendState();
                    if (sendState == AudioSendTaskStateEnum.WAIT) {
                        task = audioTask;
                        break;
                    } else if (sendState == AudioSendTaskStateEnum.SENDING) {
                        task = null;
                        break;
                    }
                } else {
                    task = null;
                    break;
                }
            }
            return task;
        }
    }

    public void addAudioTaskQueue(AudioTask audioTask) {
        synchronized (audioTaskQueue) {
            audioTaskQueue.add(audioTask);
        }
    }

    @AllArgsConstructor
    @Data
    public static class AudioTask {
        private final String taskId;
        private AudioSendTaskStateEnum sendState;
        private boolean ttsDone;

        private AudioDataDTO audioData;
    }

    public static void main(String[] args) {
        final BlockingQueue<Runnable> sendAudioTaskQueue = new LinkedBlockingQueue<>(); // 顺序任务队列

    }
}
